import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt

# 定义训练集的x
x_values = [i for i in range(11)]
x_train = np.array(x_values, dtype=np.float32)
x_train = x_train.reshape(-1, 1)
print(x_train.shape)

# 定义训练集的y
y_values = [5 * i + 1 for i in x_values]
y_train = np.array(y_values, dtype=np.float32)
y_train = y_train.reshape(-1, 1)
print(y_train.shape)


# 其实线性回归就是一个不加激活函数的全连接层

class LinearRegressionModel(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(LinearRegressionModel, self).__init__()
        self.linear = nn.Linear(input_dim, output_dim)

    def forward(self, x):
        out = self.linear(x)
        return out


def main():
    input_dim = 1
    output_dim = 1
    model = LinearRegressionModel(input_dim, output_dim)

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model.to(device)

    # 指定好参数和损失函数
    epochs = 1000
    learning_rate = 0.01
    # 优化器 SGD(Stochastic gradient descent) 随机梯度下降算法
    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
    criterion = nn.MSELoss()  # 自己定义的损失函数 均方误差损失函数

    for epoch in range(epochs):
        epoch += 1
        # 注意转行成tensor
        inputs = torch.from_numpy(x_train)
        labels = torch.from_numpy(y_train)
        # 梯度要清零每一次选代
        optimizer.zero_grad()
        # 前向传播
        outputs = model(inputs)
        # 计算损失
        loss = criterion(outputs, labels)
        # 返向传播
        loss.backward()
        # 更新权重参数
        optimizer.step()
        if epoch % 50 == 0:
            print('epoch {}, loss {}'.format(epoch, loss.item()))

    # 测试模型预测结果
    print(x_train)
    print(y_train)
    predicted = model(torch.from_numpy(x_train).requires_grad_()).data.numpy()
    print(predicted)

    x = x_train.reshape(1, -1)[0]
    y = y_train.reshape(1, -1)[0]
    y_predicted = predicted.reshape(1, -1)[0]
    plt.plot(x, y, 'o')
    plt.plot(x, y_predicted, 'o', color='red')
    plt.show()

    # torch.save(model.state_dict(), 'model.pkl')
    # model.load_state_dict(torch.load('model.pkl'))


if __name__ == '__main__':
    main()
